import numpy as np
import math
import scipy
from scipy import optimize
import control.matlab

import pat.utils as pu
#import pat.vehicles.fixed_wing.dynamic_model_python_basic as dm
import pat.vehicles.fixed_wing.dynamic_model_3d as dm

#
# equilibriums / reference trajectories
#
def get_trim_cst_path(va, gamma, dm_param, h=0., debug=False):
    """
      Find elevator, throttle and angle of attack coresponding
      to the given airspeed and and flight path
    """
    def err_func((throttle, elevator, alpha)):
        X=[0., 0., h, va, alpha, 0., 0.,  gamma+alpha, 0., 0., 0., 0.]
        U = np.zeros(dm_param.input_nb)
        U[0:dm_param.eng_nb] = throttle; U[dm_param.eng_nb+dm.i_de1] = elevator
        Xdot = dm.dyn(X, 0., U, dm_param)
        Xdot_ref = [va*math.cos(gamma), 0., -va*math.sin(gamma), 0., 0., 0., 0., 0., 0., 0., 0., 0.]
        return np.linalg.norm(Xdot - Xdot_ref)
    if debug:
        print "searching for constant path trajectory with"
        print "  va     {:f} m/s".format(va)
        print "  gamma  {:f} deg".format(pu.deg_of_rad(gamma))
    p0 = [0.2, pu.rad_of_deg(2.), pu.rad_of_deg(0.)]
    thr_e, ele_e, alpha_e = scipy.optimize.fmin_powell(err_func, p0, disp=False)
    if debug:
        print """result:
  throttle        : {:f} %
  elevator        : {:f} deg
  angle of attack : {:f} deg""".format(100.*thr_e, pu.deg_of_rad(ele_e), pu.deg_of_rad(alpha_e))
    Ue = np.zeros(dm_param.input_nb)
    Ue[0:dm_param.eng_nb] = thr_e; Ue[dm_param.eng_nb+dm.i_de1] = ele_e
    Xe = [va*math.cos(gamma), 0., va*math.sin(gamma), va, alpha_e, 0., 0., gamma+alpha_e, 0., 0., 0., 0.]
    return Xe, Ue


def get_ref_traj_cst_path(va, gamma, param, time, debug=False):
    """
    Compute a numerical straight line trajectory (aka ~equilibrium~)
    for a given flight path gammma and air speed va
    i.e.
    find throttle, elevator and angle of attack such as
    dyn() = []
    """
    Xe, Ue = get_trim_cst_path(va, gamma, dm_param, debug)
    Xe = [va*math.cos(gamma), 0., va*math.sin(gamma), va, alpha_e, 0., 0., gamma+alpha_e, 0., 0., 0., 0.]*np.ones((len(time), dm.s_size))
    Xe[:,dm.s_x] = Xe[:,dm.s_x]*time
    Ue = [thr_e, 0., ele_e, 0.]*np.ones((len(time), dm.i_size))
    return Xe, Ue



def get_trim_cst_throttle(va, throttle, dm_param, debug=False):
    """
      Find elevator, flight path and angle of attack coresponding
      to the given airspeed and throttle
    """
    def err_func((elevator, gamma, alpha)):
        X=[0., 0., 0., va, alpha, 0., 0.,  gamma+alpha, 0., 0., 0., 0.]
        U = np.zeros(dm_param.input_nb)
        U[0:dm_param.eng_nb] = throttle; U[dm_param.eng_nb+dm.i_de1] = elevator
        Xdot = dm.dyn(X, 0, U, dm_param)
        Xdot_ref = [va*math.cos(gamma), 0., -va*math.sin(gamma), 0., 0., 0., 0., 0., 0., 0., 0., 0.]
        return np.linalg.norm(Xdot - Xdot_ref)
    if debug: print """
Searching for constant path trajectory with:
  va      : {:f} m/s
  throttle: {:f} %""".format(va, throttle*100.)
    p0 = [pu.rad_of_deg(1), pu.rad_of_deg(-2.), pu.rad_of_deg(0.)]
    ele_e, gamma_e, alpha_e = scipy.optimize.fmin_powell(err_func, p0, disp=False)
    if debug: print """result
  elevator: {:f} deg
  gamma:    {:f} deg
  alpha:    {:f} deg""".format(pu.deg_of_rad(ele_e), pu.deg_of_rad(gamma_e), pu.deg_of_rad(alpha_e))
    Ue = np.zeros(dm_param.input_nb)
    Ue[0:dm_param.eng_nb] = throttle; Ue[dm_param.eng_nb+dm.i_de1] = ele_e
    Xe = [va*math.cos(gamma_e), 0., va*math.sin(gamma_e), va, alpha_e, 0., 0., gamma_e+alpha_e, 0., 0., 0., 0.]
    return Xe, Ue


def get_ref_traj_cst_throttle(va, throttle, param, time, debug=False):
    """
    Compute a numerical straight line trajectory (aka ~equilibrium~)
    for the given air speed va and throttle
    i.e.
    find elevator, flight path and angle of attack such as
    
    """
    Xe, Ue = get_trim_cst_throttle(va, throttle, dm_param, debug)
    Xe = [va*math.cos(gamma_e), va*math.sin(gamma_e), va, gamma_e, alpha_e, 0.]*np.ones((len(time), model.sv_size))
    Xe[:,model.sv_x] = Xe[:,model.sv_x]*time
    Xe[:,model.sv_h] = Xe[:,model.sv_h]*time
    Ue = [throttle, ele_e]*np.ones((len(time), model.iv_size))
    return Xe, Ue




import gobject
import xml.etree.ElementTree as ET
import StringIO

import pat.input_output.joystick   as mjs

class Reference():
    
    def __init__(self, Xr, Ur):
        self.trim = Xr, Ur
        
    def get(self, time):
        return self.trim

class step_z(Reference):

    def get(self, time):
        Xr, Ur = self.trim
        a = math.fmod(time, 16.)
        if a > 8 : Xr[dm.s_z] = -30
        else: Xr[dm.s_z] = -20
        return Xr, Ur
    

class Ap(gobject.GObject):
    """An autopilot for a fixed wing aircraft"""
    __gsignals__ = { 
        "update_fast" : (gobject.SIGNAL_RUN_LAST, gobject.TYPE_NONE,()),
        "update_slow" : (gobject.SIGNAL_RUN_LAST, gobject.TYPE_NONE,())
        } 
    def __init__(self, P_dm, ap_filename=None, js_device=None, reference=None, debug=False):
        gobject.GObject.__init__(self)
        self.P_dm = P_dm
        self.make_trim(debug=debug)
        if js_device <> None:
            self.js = mjs.Joystick(js_device)
            self.input = mjs.Joystick(js_device)
        else:
            self.js    = None
            self.input = mjs.ApInput()
        self.reference = Reference(self.Xe, self.Ue)

        self.modes=[]
        for mode_class in [ApDirect, ApAtt, ApAtt2]:#, ApAtt3]:
            self.modes.append(mode_class(P_dm, self))
        self.cur_mode = 0
        self.update_cnt = 0
        if ap_filename <> None: self.read_xml(ap_filename)
        

    def read_xml(self, filename):
        """XML parser"""
        print "ap loading xml from ",filename 
        tree = ET.parse(filename)
        root = tree.getroot()
#        self.name = root.attrib["name"]
        modes = root.find("mode")
        K = np.genfromtxt(StringIO.StringIO(modes.find("gain").text), delimiter=",")
        print K
        self.modes[1].set_gain(K)

    def set_reference(self, ref):
        self.reference = ref
        

    def make_trim(self, va=70., gamma=pu.rad_of_deg(0.), debug=False):
        """ finds trim conditions"""
        self.Xe, self.Ue = get_trim_cst_path(va, gamma, self.P_dm, debug=True)
        if debug: print "make trim: ", self.Xe, self.Ue
        self.A, self.B = pu.num_jacobian(self.Xe, self.Ue, self.P_dm, dm.dyn)

    def get_trim(self): return self.Xe, self.Ue

    def get_jacobian(self): return self.A, self.B

    def set_gain(self, K): self.modes[self.cur_mode].set_gain(K)

    def set_spec(self, omega_pitch=1., omega_roll=None, xi=0.7, debug=False):
        self.modes[self.cur_mode].set_spec(self, omega_pitch, omega_roll, xi, debug)

    def _print(self):
        self.modes[self.cur_mode]._print(self)

    def run(self, time, X):
        self.U = self.modes[self.cur_mode].run(time, X, self)
        if self.update_cnt % 10 == 0: self.emit("update_slow")
        self.update_cnt += 1
        return self.U

    def set_mode(self, no_mode):
        self.cur_mode = no_mode

    def get_mode(self):
        return self.modes[self.cur_mode]
    
    def get_U_str(self):
        thr_txt = "throttle: {:s} %".format(self.U[:self.P_dm.eng_nb]*100.)
        sfc_fmt = "  ail: {: .1f} deg   ele: {: .1f} deg   rud: {: .1f}deg"
        sfc_txt = sfc_fmt.format(pu.deg_of_rad(self.U[self.P_dm.eng_nb]),
                                 pu.deg_of_rad(self.U[self.P_dm.eng_nb+1]),
                                 pu.deg_of_rad(self.U[self.P_dm.eng_nb+2]))
        return thr_txt+sfc_txt

    def properties(self, mode=None):
        if mode == None:
            return self.modes[self.cur_mode].properties()
        else:
            return self.modes[mode].properties()


class ApDirect:
    """ A direct mode"""
    def __init__(self, P_dyn, ap):
        self.P_dyn = P_dyn
    
    def properties(self):
        return "None"

    def run(self, time, X, ap):
        Xe, Ue = ap.get_trim()
        U = np.array(Ue)
        # set engine and control surfaces
        U[0:self.P_dyn.eng_nb] += ap.input.throttle*np.ones(self.P_dyn.eng_nb)
        U[self.P_dyn.eng_nb]   += pu.rad_of_deg(30.)*ap.input.stick_roll
        U[self.P_dyn.eng_nb+1] += pu.rad_of_deg(15.)*ap.input.stick_pitch
        U[self.P_dyn.eng_nb+2] += pu.rad_of_deg(30.)*ap.input.stick_yaw
        return U
    
class LinFeedback():
    def __init__(self, P_dyn):
        self.P_dyn = P_dyn

    def set_gain(self, K): self.K = K

    def properties(self):
        return "K{:s}".format(self.K)

    def run(self, X, Xr, ap):
        dX = X-Xr
#        print "dX", dX
#        print "K", self.K
        dU = np.dot(self.K, dX)
        return dU
    
    def _print(self, ap):
        print "K\n", self.K
        Acl = ap.A + np.dot(ap.B,self.K)
        pu.print_lti_dynamics(Acl, ap.B, txt="closed loop dyanmic", print_modal_form=False)

    
class ApAtt(LinFeedback):
    """ An Attitude mode"""
    def __init__(self, P_dm, ap):
        LinFeedback.__init__(self, P_dm)
        self.set_gain(np.array([[0., 0., 0., 0., 0., 0.,  0., 0., 0. ,  0.,  0. , 0.],
                                [0., 0., 0., 0., 0., 0.,  2., 0., 0. ,  0.5, 0. , 0.],
                                [0., 0., 0., 0., 0., 0.,  0., 2., 0. ,  0.,  2.5, 0.],
                                [0., 0., 0., 0., 0.,-1.5, 0., 0., 0. ,  0.,  0. , 1.5]]))
        Xe, Ue = ap.get_trim()
        self.Xr = np.array(Xe); self.Ur = np.array(Ue)

    def properties(self):
        np.set_printoptions(precision=2, suppress=False, linewidth=140)
        return "Xr{:s}\nUr{:s}\nK\n{:s}".format(self.Xr, self.Ur, self.K)

    def run(self, time, X, ap):
        Xe, Ue = ap.get_trim()
        self.Xr = np.array(Xe); U = np.array(Ue)
        if ap.reference <> None:
            self.Xr, Ur = ap.reference.get(time)
        if ap.js <> None:
            theta_sp = -pu.rad_of_deg(30.)*ap.input.stick_pitch
            phi_sp   =  pu.rad_of_deg(30.)*ap.input.stick_roll
            self.Xr += np.array([0., 0., 0., 0., 0., 0., phi_sp, theta_sp, 0., 0., 0., 0.])
#            print theta_sp
            U[0:self.P_dyn.eng_nb] += ap.input.throttle*np.ones(self.P_dyn.eng_nb)
#        print "X", X
#        print "Xr", self.Xr
        dU = LinFeedback.run(self, X, self.Xr, ap)
#        print "dU ", dU
#        dU[0:self.P_dyn.eng_nb] = np.zeros(self.P_dyn.eng_nb)
        U += dU
#        print "U", U
        return U


class ApAtt2(ApAtt):
    def __init__(self, P_dm, ap, omega_pitch=4., omega_roll=1., xi=0.7):
        ApAtt.__init__(self, P_dm, ap)
#        self.set_spec(ap)

    def place(self, A, B, omega, xi, debug=False):
        """SISO command canonic form second order placement"""
        K = np.array([-1./B[1]*(omega**2+A[1,0]), -1./B[1]*(2.*xi*omega+A[1,1])])
        if debug:
            pu.print_lti_dynamics(A, B, txt="open loop dynamics", print_original_form=True, print_modal_form=False)
            print "K", K
            Acl = A+np.dot(B, K.transpose())
            pu.print_lti_dynamics(Acl, B, txt="obtained closed loop dynamics", print_original_form=True, print_modal_form=False)
        return K
    

    def set_spec(self, ap, omega_pitch=4., omega_roll=4., xi=0.7, debug=False):
        Xe, Ue = ap.get_trim()
        A,B = pu.num_jacobian(Xe, Ue, self.P_dyn, dm.dyn)
        if debug: pu.print_lti_dynamics(A, B, txt="set spec: open loop dyanmic",  print_original_form=True, print_modal_form=False)

        # roll
        if omega_roll <> None:
            print "roll"
            A1 = np.array([[0., 1.],[A[dm.s_p, dm.s_phi], A[dm.s_p, dm.s_p]]])
            i_ail = self.P_dyn.eng_nb
            B1 = np.array([[B[dm.s_phi, i_ail]], [B[dm.s_p, i_ail]]])
            K = self.place(A1, B1, omega_roll, xi, debug=debug)
#            K = [4, 0.5]
        else:
            K = np.zeros(2)
        self.K[dm.i_da1+self.P_dyn.eng_nb, dm.s_phi] = K[0]
        self.K[dm.i_da1+self.P_dyn.eng_nb, dm.s_p]   = K[1]

        # pitch
        if omega_pitch <> None:
            print "pitch"
            A1 = np.array([[0., 1.],[A[dm.s_q, dm.s_theta], A[dm.s_q, dm.s_q]]])
            i_ele = self.P_dyn.eng_nb+1
            B1 = np.array([[B[dm.s_theta, i_ele]], [B[dm.s_q, i_ele]]])
            K = self.place(A1, B1, omega_pitch, xi, debug=debug)
        else:
            K = np.zeros(2)
        self.K[dm.i_de1+self.P_dyn.eng_nb, dm.s_theta] = K[0]
        self.K[dm.i_de1+self.P_dyn.eng_nb, dm.s_q]     = K[1]

        # yaw
        self.K[dm.i_dr1+self.P_dyn.eng_nb, dm.s_beta] = 0.
        self.K[dm.i_dr1+self.P_dyn.eng_nb, dm.s_r]    = 0.
        
        
class ApAtt3(ApAtt):
    def __init__(self, P_dm, ap, omega_pitch=4., omega_roll=1., xi=0.7):
        ApAtt.__init__(self, P_dm, ap)
        #             z        v     alpha beta   phi  theta  psi     p       q       r
        Q = np.diag([0.01,  0.075,  0.5,   0.5,  0.5, 2.5,   2.,  0.0005, 0.0005, 0.0005])
        #            eng  ail  ele   rud
        R = np.diag([0.3, 0.1, 0.7, 0.1])
        self.set_spec(ap, Q, R)

    def set_spec(self, ap, Q, R, debug=False):
        self.Q = Q; self.R = R
        A, B = ap.get_jacobian()
        A1=A[dm.s_z:, dm.s_z:]
        B1=B[dm.s_z:,:]
        (K1, X1, E) = control.matlab.lqr(A1, B1, Q, R)
        K = np.zeros((ap.P_dm.input_nb, dm.s_size))
        K[:,dm.s_z:] = -K1
        self.K = K

        
    def properties(self):
        np.set_printoptions(precision=2, suppress=False, linewidth=140)
        return "Xr{:s}\nUr{:s}\nK\n{:s}\nQ\n{:s}\nR\n{:s}".format(self.Xr, self.Ur, self.K, np.diag(self.Q), np.diag(self.R))
